package top.zhangjianyong.blog.ai.client;

import cn.hutool.core.collection.CollectionUtil;
import jakarta.annotation.Resource;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import top.zhangjianyong.blog.ai.model.BlogData;
import top.zhangjianyong.blog.constant.ThreadPoolConstant;
import top.zhangjianyong.blog.entity.dto.BlogDTO;
import top.zhangjianyong.blog.entity.po.Category;
import top.zhangjianyong.blog.entity.po.Tag;

/**
 * @author zhangjianyong
 * @since 2025/7/8 12:47
 */
@Slf4j
@Component
public class BlogVectorStoreClient {

  @Resource private VectorStore vectorStore; // ChromaVectorStore自动注入

  // 针对1核2G服务器优化：默认搜索结果数量限制
  private static final int DEFAULT_MAX_RESULTS = 3;

  /**
   * 保存博客到向量库
   *
   * @param blogs 博客列表
   */
  public void saveBlogsToVectorStore(List<BlogData> blogs) {
    List<Document> list =
        blogs.stream()
            .map(
                blog -> {
                  log.info("保存博客：{}", blog.getId());
                  // 构造Document对象，仅用id、text、metadata
                  return new Document(blog.getId(), blog.getText(), new HashMap<>());
                })
            .toList();
    vectorStore.add(list);
  }

  public void saveBlogToVectorStore(BlogData blog) {
    saveBlogsToVectorStore(List.of(blog));
  }

  @Async(ThreadPoolConstant.ASYNC_VECTOR_STORE_EXECUTOR_BEAN)
  public void saveBlogsToVectorStoreAsync(List<BlogData> blogs) {
    saveBlogsToVectorStore(blogs);
  }

  @SuppressWarnings("unused")
  private Map<String, Object> buildMetadata(BlogDTO blogDTO) {
    Map<String, Object> metadata = new HashMap<>();
    String title = blogDTO.getTitle();
    List<Tag> tags = blogDTO.getTags();
    List<Category> categories = blogDTO.getCategories();
    if (title != null) metadata.put("title", title);
    if (tags != null) metadata.put("tag", tags);
    if (categories != null) metadata.put("category", categories);
    return metadata;
  }

  /**
   * 基于向量库检索最相关的博客
   *
   * @param query 用户问题
   * @return 相关博客的ID、标题和内容列表
   */
  public List<BlogData> searchBlogs(String query) {
    return searchBlogs(query, DEFAULT_MAX_RESULTS);
  }

  /**
   * 基于向量库检索最相关的博客（带数量限制）
   *
   * @param query 用户问题
   * @param maxResults 最大结果数量
   * @return 相关博客的ID、标题和内容列表
   */
  public List<BlogData> searchBlogs(String query, int maxResults) {
    try {
      log.info("开始向量搜索，查询: {}, 最大结果数: {}", query, maxResults);

      // 构造搜索请求
      SearchRequest request = SearchRequest.builder().query(query).similarityThreshold(0.8).build();

      // 执行搜索
      List<Document> documents = vectorStore.similaritySearch(request);

      if (CollectionUtil.isEmpty(documents)) {
        log.info("未找到相关文档");
        return new ArrayList<>();
      }

      log.info("找到{}个相关文档，限制为{}个", documents.size(), maxResults);

      // 限制结果数量
      if (documents.size() > maxResults) {
        documents = documents.subList(0, maxResults);
      }

      // 转换为BlogData对象

      return documents.stream()
          .map(
              document -> {
                BlogData blogData = new BlogData();
                blogData.setId(document.getId());
                blogData.setText(document.getText());
                return blogData;
              })
          .toList();
    } catch (OutOfMemoryError e) {
      log.error("向量搜索内存溢出，查询: {}", query, e);
      // 返回空结果，避免进一步的内存问题
      return new ArrayList<>();
    } catch (Exception e) {
      log.error("向量搜索异常，查询: {}", query, e);
      return new ArrayList<>();
    }
  }

  public void deleteBlog(List<Long> ids) {
    if (CollectionUtil.isEmpty(ids)) return;
    List<String> list = ids.stream().map(Object::toString).toList();
    vectorStore.delete(list);
  }
}
